'''
Created on Feb 2, 2010

@author: oabalbin
'''

#import sys;
#import string;
#import StringIO;
import os;
import sys;
import glob;
import MySQLdb;
import numpy as np
from optparse import OptionParser
from collections import deque

# My libraries
import signatures.parsers.read_microarray_data as pm
import signatures.parsers.read_gene_lists as gpm
import signatures.db.tables as dt
import signatures.db.query as dq
import signatures.common.classes as cm #record, arrayMatrix
import signatures.preprocess.tools as tl #filter_constant_genes, get_variable_genes, 
import signatures.seedgenes.seedlist as sd #get_seed_genes, write_varin_file
import signatures.bfrm.paramfile as bp
import signatures.bfrm.bfrm_run as br

class pipeline():
    
    def __init__(self, dbhost, dbuser, dbpasswd, database):
        
        self.mt = dt.tables(dbhost, dbuser, dbpasswd, database)
        self.mq = dq.query(dbhost, dbuser, dbpasswd, database)
        #table information
        self.tableDescription = 'probeID VARCHAR(40), geneSym VARCHAR(40), sample VARCHAR(40), log_fold FLOAT(20,14), INDEX probe (probeID), INDEX gene (geneSym), INDEX sample (sample)'
        self.tableName = 'NCI60_ERG_2'
        #self.tableName = 'erg_signatures_only'
        # Parameters for pre-processing data:
        # percentage of low variable gene to be filter out of the analysis
        self.lvar=25
        # percentage of low median gene to be filter out of the analysis
        self.lmed=25
        
        # sample classes
        #self.spclass0=deque(['RWPE_3_RWPE_7','RWPE_7_CPP','CPP_RWPE19','RWPE_3_CPP','RWPE_CPP_dup','CPP_RWPE_7'])
        
        self.spclass0=deque(['RWPE_3_RWPE_7','RWPE_7_CPP','CPP_RWPE19','RWPE_3_CPP','RWPE_CPP_dup','CPP_RWPE_7', 'Vcap_SiRNA_ERG_nontarget','Vcap_SiRNA_ERG_nontarget_dup','Vcap_SiRNA_ERG_nontarget_swap',
                       'Vcap_SiRNA_ERG_nontarget_swap_dup','LNCAP_SiRNAETV1_notarget1','LNCAP_SiRNAETV1_notarget2','LNCAP_SiRNAETV1_notarget1_swap','LNCAP_SiRNAETV1_notarget2_swap'])
        
        
        self.spclass1=deque(['PREC_AD_ERG1_LACZ','RWPE_L_ERG1_GUS','PREC_AD_ERG1_LACZ_dup','PREC_AD_ERG1_LACZ_dyeswap','RWPE_L_ERG1_GUS_dyes_dup',
                      'RWPE_L_ERG1_GUS_dyeswap', 'RWPE_L_ERG1_GUS_dup','PREC_AD_ERG1_LACZ_dyeswap_dup','PREC_AD_ETV1_LACZ','PREC_AD_ETV1_LACZ_dyeswap',
                      'RWPE_L_ETV1_GUS_dyeswap','RWPE_L_ETV1_GUS','PREC_AD_ETV1_LACZ_dup','RWPE_L_ETV1_GUS_dup','PREC_AD_ETV1_LACZ_dyeswap_dup','RWPE_L_ETV1_GUS_dyeswap_dup'])
        
        #self.spcltest=[]

        '''
        self.spcltest= deque(['Vcap_SiRNA_ERG_nontarget','Vcap_SiRNA_ERG_nontarget_dup','Vcap_SiRNA_ERG_nontarget_swap',
                       'Vcap_SiRNA_ERG_nontarget_swap_dup','LNCAP_SiRNAETV1_notarget1','LNCAP_SiRNAETV1_notarget2','LNCAP_SiRNAETV1_notarget1_swap','LNCAP_SiRNAETV1_notarget2_swap',
                       'RWPE_L_ETV1_GUS','PREC_AD_ETV1_LACZ','22RV1','22RV1_dup','RWPE_1','RWPE_1_dup','LnCap','LnCap_dup',])
        
        
        self.spcltest=['GEaN23','GEaN25','GEaN27','GEaN29','GEaN31','GEaN32','GEaN33','GEaT20','GEaT32','GEaT38','GEaT40','GEaT42','GEaT43','GEaT45','GEaT47',
                       'GEaT49','GEaT52','GEaT53','GEaT54','GEaT55','GEaT56','GEaT57','GEaT58','GEaT61','GEaT62','GEaT64','GEaT65','GEaT66','GEaT76','GEaT77','GEaT78','GEaT79',
                       'GEaT80','GEaT81','GEaT82','GEaT83','GEaT84','GEaT85','GEaM15','GEaM16','GEaM20','GEaM21','GEaM22','GEaM23','GEaM26','GEaM28','GEaM29','GEaM33','GEaM34','GEaM36',
                       'GEaM37','GEaM38','GEaM39','GEaM40','GEaM41','GEaM43','GEaM44']
        
        '''
             
        self.spcltest=deque(['MCF7', 'MDA-MB-231/ATCC', 'HS 578T', 'BT-549', 'T-47D', 'SF-268', 'SF-295', 'SF-539', 'SNB-19', 'SNB-75', 'U251', 'COLO 205', 'HCC-2998', 'HCT-116', 
                       'HCT-15', 'HT29', 'KM12', 'SW-620', 'CCRF-CEM', 'HL-60(TB)', 'K-562', 'MOLT-4', 'RPMI-8226', 'SR', 'LOX IMVI', 'MALME-3M', 'M14', 'SK-MEL-2', 'SK-MEL-28', 
                       'SK-MEL-5', 'UACC-257', 'UACC-62', 'MDA-MB-435', 'MDA-N', 'A549/ATCC', 'EKVX', 'HOP-62', 'HOP-92', 'NCI-H226', 'NCI-H23', 'NCI-H322M', 'NCI-H460', 'NCI-H522', 
                       'IGROV1', 'OVCAR-3', 'OVCAR-4', 'OVCAR-5', 'OVCAR-8', 'SK-OV-3', 'NCI/ADR-RES', 'PC-3', 'DU-145', '786-0', 'A498', 'ACHN', 'CAKI-1', 'RXF 393', 'SN12C', 'TK-10', 
                       'UO-31'])
        '''
        self.spclass1=['PREC_AD_ERG1_LACZ','RWPE_L_ERG1_GUS','PREC_AD_ERG1_LACZ_dup','PREC_AD_ERG1_LACZ_dyeswap',
                  'RWPE_L_ERG1_GUS_dyes_dup','RWPE_L_ERG1_GUS_dyeswap','PC3_ERG_dETS_1', 'PREC_RAM_ERG', 
                  'HPAEpiC_RAM_ERG', 'HMEpiC_RAM_ERG']
        self.spcltest=['RWPE_3_RWPE_7','RWPE_CPP_dup','RWPE_L_ERG1_GUS_dup','PC3_ERG_dETS_3','PREC_AD_ERG1_LACZ_dyeswap_dup',
                    'PREC_RAM_ERG_dup','HPAEpiC_RAM_ERG_dup','HMEpiC_RAM_ERG_dup']
        
    
        self.spclass1=['PREC_AD_ERG1_LACZ','RWPE_L_ERG1_GUS','PREC_AD_ERG1_LACZ_dup','PREC_AD_ERG1_LACZ_dyeswap','RWPE_L_ERG1_GUS_dyes_dup',
                      'RWPE_L_ERG1_GUS_dyeswap','PC3_ERG_dETS_1', 'PREC_RAM_ERG', 'HPAEpiC_RAM_ERG', 'HMEpiC_RAM_ERG', 'RWPE_L_ERG1_GUS_dup',
                      'PC3_ERG_dETS_3','PREC_AD_ERG1_LACZ_dyeswap_dup','PREC_RAM_ERG_dup','HPAEpiC_RAM_ERG_dup','HMEpiC_RAM_ERG_dup']
        
        '''        
        '''
        self.spcltest=['RWPE_3_RWPE_7','RWPE_7_CPP','CPP_RWPE19','RWPE_3_CPP','RWPE_CPP_dup','CPP_RWPE_7','PREC_AD_ETV1_LACZ', 
                      'PREC_AD_ETV1_LACZ_dyeswap','RWPE_L_ETV1_GUS_dyeswap','RWPE_L_ETV1_GUS','PREC_AD_ETV1_LACZ_dup','RWPE_L_ETV1_GUS_dup',
                     'PREC_AD_ETV1_LACZ_dyeswap_dup','RWPE_L_ETV1_GUS_dyeswap_dup']
        
        self.spcltest=[]
        '''
        
    
    
    def parse_input_annotfiles(self, annotInputFile):
    
        expr_file, seed_file, outdirs, db_file, samples, bfrm, typefile={},{},{},{},{},{},'nofile'
    
        for line in annotInputFile:
            fields = line.strip('\n')
                 
            if fields[0][0] == '':
                continue
            
            if fields[0][0] == '#':
                typefile = fields
                continue
            
            fields = fields.split('=')
            dict = {'#Array expression':expr_file, '#Seed genes':seed_file,'#Database file':db_file,'#Samples':samples,'#Output Directories':outdirs,\
                    '#BFRM':bfrm}[typefile]
            dict[fields[0]] = fields[1]
            
        return expr_file, seed_file, db_file, samples, bfrm, outdirs
    

    def read_files_folder(self,folderpath,ext):
        ''' '''
        # Read files in folder
        myfiles=[]
        for infile in glob.glob( os.path.join(folderpath, '*'+ext) ):
            myfiles.append(infile)
            #print "current file is: " + infile
        return myfiles
    
    
    def process_microarray_files(self,exp_files,tp):
        dbf=[]
        for f in exp_files:
            of=f+'_dbfile'
            tp.microarray_exp_file(open(f),open(of,'w'))
            dbf.append(of)
        
        return dbf
    
    def create_database(self,dbfiles,doit=False):
        """
        It creates a databases if doit = True
        """
        if doit:
            print 'Making tables'            
            for dbf in dbfiles:     
                if not(self.mt.table_existence(self.tableName)):
                    self.mt.create_table(self.tableName,self.tableDescription)
                    
                self.mt.load_data_table(self.tableName,dbf)
                
    
    def filter_expMat(self, samplelist, initgenelist, tableName, seedfile, dofilter=False):     
        '''
        Returns an arrayMatrix with the final expression matrix.
        From the matrix you: 
        
        a) filter the all genes to remove noise, keep only genes var >= lvar, med >= lmed
        b) if another optional list is provided, you want to see
           b.1) how many of the new genes are variable. 
           b.2) Produce a expression matrix with only genes in b.1
        
        '''
        self.mt.create_connector()
        self.mq.create_connector()
        # Expression array matrix.            
        expMat = self.mq.get_exp_mat_2(samplelist, initgenelist, self.tableName)
        # Normalize the expression values data, using oncomine approach.
        expMat.normalization()
        
        if dofilter:
            # Find variable genes and modify he expMat.
            genelist = tl.filter_out_constant_genes(expMat, self.lvar, self.lmed, False)       
            # Get the index and name of the variable genes for a particular experiment
            mgenelist, dictmgl =  tl.get_variable_genes(expMat,genelist)
            
            expMat = tl.get_filtered_expMat(expMat,dictmgl,len(samplelist))
            expMat.create_gen_index()
     
        # list of external genes to be used as the "universe of genes"
        # e.i: this could be the ERG interactome
        if seedfile:
            extgenlist = tp.list_of_names_in_line(open(seedfile))
            # Index of genes belonging to the "universe" which also have var .=lvar and med >=lmedS
            #print 'Genes in matrix expMat'
            #print expMat.genInd
            
            dictmgl2, genenotfound = sd.get_gen_index(expMat,extgenlist)
            #mgenelist, dictmgl =  tl.get_variable_genes(expMat,finalgenes)
            #print genenotfound
            
            #print 'Genes in present in the universe list'
            #print dictmgl2 
            expMat = tl.get_filtered_expMat(expMat,dictmgl2,len(samplelist))
        
        #expMat.normalization()
        expMat.log2_norm()
        #print mgenelist
        #print 'Genes not found in the variable genes'        
        #print genenotfound
        #print 'Number of genes ' + str(genenotfound)
        
        return expMat 


##########


"""
process the input data. And make the database

Handle different intentions:
    a) generate files for BFRM
    b) generate structures for running BRFM
    c) generate files for 
    f) EYE!!! need to be able to subquery the big database using a sample list.
    c.2) Allowed for outputing the median of gene, or all probe sets for a gen
        Consider the case when different platforms are used.
        Remember here to modify the function get_exp_mat_2 and the get_genelist_plain in query
    d) allowed to output the raw or the normalized data
    e) handle different classes of samples input for the user
    f) EYE!!! need to be able to subquery the big database using a sample list.

"""





if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--annotFile", dest="annotFile",
                            help="annotation file for all files to use")
    optionparser.add_option("-d", action="store_true", dest="dbtrue",
                            help="create a database. Bool")

    
    (options, args) = optionparser.parse_args()
       
    # start the pipeline. 
    # open the database
    pipe = pipeline("localhost", "oabalbin", "oscar", "signatures")
    expr_folder, seed_file, db_file, samples, bfrm, outdirs = pipe.parse_input_annotfiles(open(options.annotFile))    
    #read folder for expression profiles
    exp_files = pipe.read_files_folder(expr_folder['expdir'],'.txt')
    
    ####### output files
    
    outfile1 = outdirs['outdir']+'bfrm/dataset_mat.txt'  # Expression values
    outfile2 = outdirs['outdir']+'bfrm/dataset_gen.txt'  # geneNames
    outfile3 = outdirs['outdir']+'bfrm/varin.txt'        # seed genes
    outfile4 = outdirs['outdir']+'bfrm/parameters.txt'   # bfrm parameters file
    outfile5 = outdirs['outdir']+'svd/mat_exp_SiRNA_NCI60.txt'   # bfrm parameters file
    outfile6 = outdirs['outdir']+'svd/gen_exp_SiRNA_NCI60.txt'   # bfrm parameters file

    ######### Get raw data  ############
    # Get the database files
    tp = pm.parser()
    #dbfiles = pipe.process_microarray_files(exp_files,tp)
    pipe.mt.create_connector()
    
    gtp = gpm.geneparser()
    ######## Create database
    #pipe.create_database(dbfiles,options.dbtrue)
    
    ########## Make queries ############
    
    pipe.mq.create_connector()
    # list of samples:
    #samplelist = tp.list_of_names_in_line(open(samples['samples']))
    samplelist = tp.list_of_names_in_line(open(samples['samples']))
    # list of all ucsc potential genes to include in study
    #initgenelist = gtp.list_of_genes_kgXref(open(seed_file['allgene']))
    #initgenelist = tp.list_of_names_in_line(open(seed_file['univfile']))
    initgenelist = tp.list_of_names(open(seed_file['univfile']))    
    # Do query to obtain the expression matrix 
    
    #extgenlist = tp.list_of_names_in_line(open(seed_file['allgene']))
    
    # print initgenelist
    print samplelist
    expMat = pipe.mq.get_exp_mat_2(samplelist, initgenelist, pipe.tableName)
    # Normalize the expression values data
    expMat.normalization()
    #expMat.log2_norm()

    ######### Pre-process Raw Data ###############
    # NOTE: This parameter pipe.lvar=25 can be modified to control variability filter
    # Get the index of genes with var >=lvar, and med >=lmed 
    
    #expMat = pipe.filter_expMat(samplelist, initgenelist, pipe.tableName, '', dofilter=False) #seed_file['univfile']
    
    #print expMat_bfrm.geneSynd
    #print str(len(finalgenes)), str(len(genenotfound))
        
    ########## Get seed genes
    #seedgenes = tp.list_of_names(open(seed_file['seedfile']))
    #seedgenes = tp.list_of_names_in_line(open(seed_file['seedfile']))
    #seeding, seednotfound = sd.get_gen_index(expMat_bfrm,seedgenes)
    #seeding, seednotfound = sd.get_seed_genes(finalgenes,seedgenes)
    #print len(seedgenes)
    ######## Write expression matrix for variable genes
    ######## write in varin.txt file the index of seed of genes. To use in BRFM
    
    ######### BFRM
    # Choose as seed all genes in the universe
    #sd.write_varin_file(expMat_bfrm.genInd.values(),open(outfile3,'w'))
    #pipe.mq.write_variable_genes(expMat_bfrm, expMat_bfrm.genInd.values(), open(outfile1,'w'), open(outfile2,'w'))
    
    # Choose as seed the genes given by the seedfile
    #sd.write_varin_file(seeding.values(),open(outfile3,'w'))
    #pipe.mq.write_variable_genes_nfact(expMat_bfrm, seeding.values(), open(outfile1,'w'), open(outfile2,'w'), True)
    
    
    
    ########### SDV
    # write output for the sdv program
    # Note 02-15-2010: The sample class are going to be hard code for now. Latter on can be readed from the sample file
    #pipe.mq.write_variable_genes_sdv(expMat_bfrm, expMat_bfrm.genInd.values(), open(outfile5,'w'), open(outfile6,'w'), pipe.spclass0, pipe.spclass1, pipe.spcltest)
    
    # For the sdv program when using all genes  summarized by median of gene expression
    #pipe.mq.write_variable_genes_sdv(expMat, expMat.genInd.values(), open(outfile5,'w'), open(outfile6,'w'), pipe.spclass0, pipe.spclass1, pipe.spcltest) 
    forallProbes=range((expMat.expVal.shape[0]))
    pipe.mq.write_variable_genes_sdv(expMat, forallProbes, open(outfile5,'w'), open(outfile6,'w'), pipe.spclass0, pipe.spclass1, pipe.spcltest)
    sys.exit(0)
    ######### RUN BFRM #######
    ######### Create the parameters file
    # Call to parameters function
    '''
    parmeter_val(self, NLatentFactors, NObservations, NVariables, DataFile, EvolVarIn, EvolVarInFile,
                  EvolIncludeVariableThreshold=0.75, EvolIncludeFactorThreshold=0.75,
                  EvolMaximumFactors=10, EvolMaximumVariables=150):
    '''
    NLatentFactors = 11
    NObservations = len(samplelist)
    NVariables = len(finalgenes)
    DataFile = outfile1
    EvolVarIn = len(seeding)
    EvolVarInFile = outfile3
    EvolIncludeVariableThreshold=0.75
    EvolIncludeFactorThreshold=0.75
    EvolMaximumFactors=20
    EvolMaximumVariables=200
    
    print EvolVarIn
    print EvolMaximumVariables

    
    bfbp = bp.paramfile()
    
    bfbp.parmeter_val(NLatentFactors, NObservations, NVariables, DataFile, EvolVarIn, EvolVarInFile, 
                      EvolIncludeVariableThreshold, EvolIncludeFactorThreshold, EvolMaximumFactors, EvolMaximumVariables)
    
    bfbp.write_bfrmparam_file(open(outfile4,'w'))
    
    #sys.exit(0)
    ########## Create an bfrm_run object 
    bfrm = br.bfrm_run(bfrm['bfrm_p'],bfrm['bfrm'])
    bfrm.execute_bfrm()
    sys.exit(0)